from geesibling.core.types import Graph, Node
from geesibling.core.lib._graph import DataType,DeviceType,Device,search_policy
from geesibling.tools import log
from geesibling.adapters.pytorch.getTorchGraph import is_tensor,getTorchGraph
from torchvision import models
import torch
from torch import fx
import re

__doc__ = """
convert torch.fx graph to Geesibling Graph
Author: jirongting
datetime: 2024.05.17
version: 3 2024.05.17 second commit
"""

_OP_NAME=[
    'transformer','lm_head', #gpt2
    'bert','classifier', # bert
    'model', # llama2
]

_MODEL_PATTERN={
    'GPT2': r"(transformer)_(?:h_)?(\d+)",
    'GPT2ForSequenceClassification': r"(transformer)_(?:h_)?(\d+)",
    'BertForSequenceClassification': r"(bert)_[^_]*_(\d+)",
    'Llama2': r"(model)_[^_]*_(\d+)",
}

"""
add inputs for each geesibling node

Parameters:
 node - geesibling Node
 fx_node - torch.fx Node

"""
def add_inputs(node:Node, fx_node:fx.Node):
    for input_node in fx_node._input_nodes:
        node.add_input(input_node.name)


"""
add outputs for each geesibling node

Parameters:
 node - geesibling Node
 fx_node - torch.fx Node

"""
def add_outputs(node:Node, fx_node:fx.Node):
    for output_node in fx_node.users:
        node.add_output(output_node.name)

"""
add input_port for each geesibling node

Parameters:
 node - geesibling Node
 fx_node - torch.fx Node

"""

def add_abstract_inputs(node:Node,fx_node:fx.Node):
    for i, input_node in enumerate(fx_node._input_nodes):
        if is_tensor(input_node.meta): # if the input node has output, add the shape into inputport else [] 
            node.add_inputport(input_node.name,0,i,DataType.F32,list(input_node.meta['tensor_meta'].shape))
        else:
            node.add_inputport(input_node.name,0,i,DataType.F32,[])

"""
add output_port for each geesibling node

Parameters:
 node - geesibling Node
 fx_node - torch.fx Node

"""
def add_abstract_outputs(node:Node,fx_node:fx.Node):
    if fx_node.name=="output": # if is the last node
        #print(type(fx_node.meta["tensor_meta"]))
        shape_dict = fx_node.meta["tensor_meta"]
        output_shape = shape_dict.get(next(iter(shape_dict)))
        output_shape= output_shape[0]
        node.add_outputport(DataType.F32,list(output_shape),0)
    elif is_tensor(fx_node.meta): # if the input node has output, add the shape into inputport else []:
        node.add_outputport(DataType.F32,list(fx_node.meta['tensor_meta'].shape),0)
    else:
        node.add_outputport(DataType.F32,[],0)


"""
add op group for transformers models

Parameters:
 node - geesibling Node
 fx_node - torch.fx Node

"""
def add_op_group(node:Node,fx_node:fx.Node,model_name):
    pattern = _MODEL_PATTERN[model_name]
    if any(op_name in fx_node.name for op_name in _OP_NAME):
        if re.match(pattern=pattern,string=fx_node.name):
            group_name = re.search(pattern,fx_node.name).group(0)
        else:
            group_name = fx_node.name
        node.attrs={"colocation_group":group_name}

"""
fx graph convert to Geesibling Graph

Parameters:
 traced - transformers.utils.fx graph

Returns:
 return a Geesibling Graph

"""
def graphToGeeGraph(traced:fx.graph_module.GraphModule,model_name):
    geesiblingGraph = Graph()
    for fx_node in traced.graph.nodes:
        if any(op_name in fx_node.name for op_name in _OP_NAME):
            node = Node(fx_node.name,fx_node.op) # attr name,op
            add_inputs(node,fx_node) # attr inputs
            add_outputs(node,fx_node)# attr outputs
            add_abstract_inputs(node,fx_node)
            add_abstract_outputs(node,fx_node)
            add_op_group(node,fx_node,model_name=model_name)
            geesiblingGraph.add_node(node)

    return geesiblingGraph